package com.idp.web.filter;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

public class ResponseWrapper extends HttpServletResponseWrapper {

	private ByteArrayOutputStream byteArrayOutputStream;
    private ServletOutputStream servletOutputStream;
    private PrintWriter writer;

	public ResponseWrapper(HttpServletResponse response) {
		super(response);
		byteArrayOutputStream=new ByteArrayOutputStream();
        servletOutputStream=new WapperedOutputStream(byteArrayOutputStream);
        writer = new PrintWriter(new OutputStreamWriter(byteArrayOutputStream,StandardCharsets.UTF_8));
	}

	@Override
    public ServletOutputStream getOutputStream(){
        return servletOutputStream;
    }

    @Override
    public PrintWriter getWriter() {
        return writer;
    }

    @Override
    public void flushBuffer(){
        if (servletOutputStream!=null){
            try {
                servletOutputStream.flush();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        if (writer!=null) {
            writer.flush();
        }
    }

    @Override
    public void reset() {
        byteArrayOutputStream.reset();
    }

    public String getResponseData(String charset) throws IOException {
        flushBuffer();
        byte[] bytes = byteArrayOutputStream.toByteArray();

        return new String(bytes,charset);
    }

	private class WapperedOutputStream extends ServletOutputStream{

        private ByteArrayOutputStream bos;

        public WapperedOutputStream(ByteArrayOutputStream stream){
            bos = stream;
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {

        }

        @Override
        public void write(int b){
            bos.write(b);
        }
    }

	public byte[] getContent(){
        flushBuffer();

        return byteArrayOutputStream.toByteArray();
    }
}
